import argparse
import os
import zipfile
import PyPDF2
import py7zr


def read_passwords_from_file(passwords_file_path):
    """
    从文本文件中读取密码列表。

    :param passwords_file_path: 包含密码的文本文件路径
    :return: 包含密码的列表
    """
    try:
        with open(passwords_file_path, 'r', encoding='utf-8') as file:
            passwords = [line.strip() for line in file.readlines()]
        return passwords
    except FileNotFoundError:
        print(f"文件 {passwords_file_path} 不存在。")


def zip_with_password(zip_file_path, password):
    """
    破解密码解压zip文件。

    :param zip_file_path: ZIP文件路径
    :param password: <PASSWORD>
    """
    if type(password) == list:
        for p in password:
            try:
                with zipfile.ZipFile(zip_file_path, 'r', zipfile.ZIP_STORED) as z:
                    z.extractall("./", pwd=p.encode("utf-8"))
                print(f"密码正确: {p}")
                break
            except:
                continue
        else:
            print("使用-t参数指定密码文件中的密码不匹配")
    elif password is None:
        try:
            with zipfile.ZipFile(zip_file_path, 'r', zipfile.ZIP_STORED) as z:
                z.extractall("./")
            print("解压成功!")
        except Exception as e:
            print(f"需要密码，使用-p参数指定密码或者使用-t参数指定密码文件，报错：{e}")

    else:
        try:
            with zipfile.ZipFile(zip_file_path, 'r', zipfile.ZIP_STORED) as z:
                z.extractall("./", pwd=password.encode("utf-8"))
                print("解压成功!")
        except Exception as e:
            print(f"密码报错,使用-p参数指定密码或者使用-t参数指定密码文件，报错:{e}")


def decompress_7z(file_path, password):
    """
    解压7z文件。
    :param file_path: 7z文件的路径
    :param password: <PASSWORD>
    :return: None
    """
    if type(password) == list:
        for p in password:
            try:
                with py7zr.SevenZipFile(file_path, mode='r', password=p) as archive:
                    archive.extractall("./")
                print(f"解压成功，密码为：{p}")
                break
            except:
                continue
        else:
            print("使用-t参数指定密码文件中的密码不匹配")
    elif password is None:
        try:
            with py7zr.SevenZipFile(file_path, mode='r') as archive:
                archive.extractall("./")
            print(f"解压成功!")
        except Exception as e:
            print(f"解压失败，请检查文件是否损害！报错结果:{e}")
    else:
        try:
            with py7zr.SevenZipFile(file_path, mode='r', password=password) as archive:
                archive.extractall("./")
                print("解压成功！")
        except Exception as e:
            print(f"密码错误，请检查密码是否正确！报错结果:{e}")


def decrypt_pdf(input_path, passwords):
    """
    尝试用给定的密码列表解密PDF文件。

    :param input_path: 需要解密的PDF文件路径
    :param passwords: 密码列表或密码
    """
    if type(passwords) == list:
        for password in passwords:
            try:
                with open(input_path, 'rb') as file:
                    pdf_reader = PyPDF2.PdfReader(file)
                    if pdf_reader.is_encrypted:
                        pdf_reader.decrypt(password)

                        decrypted_file_path = "new_" + input_path
                        with open(decrypted_file_path, 'wb') as output_file:
                            writer = PyPDF2.PdfWriter()
                            for page_num in pdf_reader.pages:
                                writer.add_page(page_num)
                            writer.write(output_file)

                        print(f"成功解密，密码为：{password}")
                        break
            except Exception as e:
                print(f"使用密码'{password}'时出错：{e}")
        else:
            print("使用-t参数指定密码文件中的密码不匹配")
    elif passwords is None:
        try:
            with open(input_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                if pdf_reader.is_encrypted:
                    print(f"文件{input_path}需要密码!")
                else:
                    decrypted_file_path = "new_" + input_path
                    with open(decrypted_file_path, 'wb') as output_file:
                        writer = PyPDF2.PdfWriter()
                        for page_num in pdf_reader.pages:
                            writer.add_page(page_num)
                        writer.write(output_file)
        except Exception as e:
            print(f"打开文件时出错：{e}")
    else:
        with open(input_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            if pdf_reader.is_encrypted:
                pdf_reader.decrypt(passwords)

                decrypted_file_path = "new_" + input_path
                with open(decrypted_file_path, 'wb') as output_file:
                    writer = PyPDF2.PdfWriter()
                    for page_num in pdf_reader.pages:
                        writer.add_page(page_num)
                    writer.write(output_file)


def determine_action(file_path):
    _, ext = os.path.splitext(file_path)
    ext = ext.lower()

    if ext == ".pdf":
        return "解密PDF"
    elif ext == ".7z":
        return "解压7z"
    elif ext in {".zip", ".rar"}:
        return "加密zip"
    else:
        raise ValueError("不支持的文件类型：{}".format(ext))


def main():
    parser = argparse.ArgumentParser(description="处理PDF、7z和zip文件")

    parser.add_argument("file_path", help="待处理文件的路径")

    parser.add_argument("-t", "--passwords-file", default="passwords.txt",
                        help="包含密码的文本文件路径，每行一个密码。默认：passwords.txt")

    parser.add_argument("-o", "--no-password", action="store_true",
                        help="指示文件无需密码")

    parser.add_argument("-p", "--password", default=None,
                        help="明确指定单个密码")

    args = parser.parse_args()

    action = determine_action(args.file_path)

    if args.no_password:
        password = None
    elif args.password:
        password = args.password
    else:
        password = read_passwords_from_file(args.passwords_file)

    if action == "解密PDF":
        decrypt_pdf(args.file_path, password)
    elif action == "解压7z":
        decompress_7z(args.file_path, password)
    elif action == "加密zip":
        zip_with_password(args.file_path, password)
    else:
        raise ValueError("不支持的操作：{}".format(action))


if __name__ == "__main__":
    main()
